package track

import (
	"github.com/bluenviron/gortsplib/v4/pkg/rtcpreceiver"
	"github.com/bluenviron/gortsplib/v4/pkg/rtcpsender"
	"github.com/bluenviron/gortsplib/v4/pkg/rtplossdetector"
	"github.com/bluenviron/gortsplib/v4/pkg/rtpreorderer"
	"github.com/bluenviron/gortsplib/v4/pkg/rtptime"
	"log"
	"sync"
	"time"

	"github.com/pion/rtcp"
	"github.com/pion/rtp"
)

type (
	WriteRTPFunc    func(t Track, packet *rtp.Packet) error
	WriteRTCPFunc   func(t Track, packets []rtcp.Packet) error
	PacketPTSFunc   func(t Track, pkt *rtp.Packet) (time.Duration, bool)
	SessionTimeFunc func() (time.Time, bool)

	OnNewTrackFunc func(t Track)
	OnReadRTPFunc  func(t Track, packet *rtp.Packet)
	OnReadRTCPFunc func(t Track, packets []rtcp.Packet)
	OnFrameFunc    func(t Track, frame *Frame)

	OnLossRTPFunc func(t Track, sequenceNumber uint16, diff int)
)

type Frame struct {
	Codec CodecId
	PTS   time.Duration
	NTP   time.Time
	Data  []byte
}

type Parma struct {
	_parent Track

	Id         int
	FmtContext FormatContext

	InitialTimestamp      uint32
	InitialSequenceNumber uint16
	PayloadMaxSize        int

	// 由外部提供实现, 内部调用
	_WriteRTP    WriteRTPFunc    // 发送rtp
	_WriteRTCP   WriteRTCPFunc   // 发送rtcp
	_PacketPTS   PacketPTSFunc   // 计算packet的pts
	_SessionTime SessionTimeFunc //

	// 由外部调用, 内部转发给订阅者, 中间进行一次拦截(处理包)
	_OnReadRTCP OnReadRTCPFunc // 收到rtcp packet
	_OnReadRTP  OnReadRTPFunc  // 收到rtp packet
	_OnFrame    OnFrameFunc    // 收到frame
	_OnLossRTP  OnLossRTPFunc  // rtp丢失

	rtpTimeEncoder *rtptime.Encoder

	rtpLossDetector *rtplossdetector.LossDetector
	rtpReorder      *rtpreorderer.Reorderer

	rtcpReceiver *rtcpreceiver.RTCPReceiver
	rtcpSender   *rtcpsender.RTCPSender
	mux          sync.Mutex
}

func newParma(parent Track) *Parma {
	p := &Parma{
		_parent: parent,
		Id:      NewID(),
		FmtContext: FormatContext{
			MediaType: MediaTypeApplication,
			Direction: DirectionRecvOnly,
			MediaInfo: MediaInfo{
				RtpMap:     RtpMapInfo{},
				Fmtp:       FmtpInfo{},
				SSRC:       NewSSRC(),
				MediaProto: []string{"RTP", "AVP"},
			},
			SessionInfo: SessionInfo{},
		},
		InitialTimestamp:      0,
		InitialSequenceNumber: 0,
		PayloadMaxSize:        RTPMaxPayloadSize,
		rtpLossDetector:       rtplossdetector.New(),
		rtpReorder:            rtpreorderer.New(),
		rtcpReceiver:          nil,
		rtcpSender:            nil,
	}

	return p
}

func (p *Parma) open() {
	p.mux.Lock()
	defer p.mux.Unlock()

	p.rtpTimeEncoder = rtptime.NewEncoder(p._parent.ClockRate(), p.InitialTimestamp)

	p.rtcpReceiver, _ = rtcpreceiver.New(
		p._parent.ClockRate(), nil, time.Minute, time.Now,
		func(packet rtcp.Packet) {
			_ = p._parent.WriteRTCP([]rtcp.Packet{packet})
		},
	)
	p.rtcpSender = rtcpsender.New(p._parent.ClockRate(), time.Minute, time.Now,
		func(packet rtcp.Packet) {
			_ = p._parent.WriteRTCP([]rtcp.Packet{packet})
		},
	)

}

func (p *Parma) close() {
	p.mux.Lock()
	defer p.mux.Unlock()

	if p.rtcpReceiver != nil {
		p.rtcpReceiver.Close()
	}
	if p.rtcpSender != nil {
		p.rtcpSender.Close()
	}
}

func (p *Parma) SetDirection(direction Direction) *Parma {
	p.FmtContext.Direction = direction
	return p
}

func (p *Parma) SetCodec(codec CodecId) *Parma {
	p.FmtContext.GetRtpMap().Codec = codec
	return p
}

func (p *Parma) SetSSRC(ssrc uint32) *Parma {
	p.FmtContext.MediaInfo.SSRC = ssrc
	return p
}

func (p *Parma) SetPayloadType(pt uint8) *Parma {
	p.FmtContext.GetRtpMap().PayloadType = pt
	return p
}

func (p *Parma) SetClockRate(clockRate int) *Parma {
	p.FmtContext.GetRtpMap().ClockRate = clockRate
	return p
}

func (p *Parma) SetChannel(channel int) *Parma {
	p.FmtContext.GetRtpMap().Parameters = channel
	return p
}

func (p *Parma) WriteRTP(t Track, packet *rtp.Packet) error {
	p.rtcpSender.ProcessPacket(packet, time.Now(), t.PTSEqualsDTS(packet))
	if p._WriteRTP != nil {
		return p._WriteRTP(t, packet)
	}
	return nil
}
func (p *Parma) SetWriteRTPFunc(fn WriteRTPFunc) *Parma {
	p._WriteRTP = fn
	return p
}

func (p *Parma) WriteRTCP(t Track, packets []rtcp.Packet) error {
	if p._WriteRTCP != nil {
		return p._WriteRTCP(t, packets)
	}
	return nil
}
func (p *Parma) SetWriteRTCPFunc(fn WriteRTCPFunc) *Parma {
	p._WriteRTCP = fn
	return p
}

func (p *Parma) SessionTime() (time.Time, bool) {
	if p._SessionTime != nil {
		return p._SessionTime()
	}
	return time.Time{}, false
}

func (p *Parma) PacketPTS(t Track, pkt *rtp.Packet) (time.Duration, bool) {
	if p._PacketPTS != nil {
		return p._PacketPTS(t, pkt)
	}
	return 0, false
}
func (p *Parma) SetPacketPTSFunc(fn PacketPTSFunc) *Parma {
	p._PacketPTS = fn
	return p
}

func (p *Parma) PacketNTP(pkt *rtp.Packet) (time.Time, bool) {
	if p.rtcpReceiver == nil {
		return time.Time{}, false
	}

	return p.rtcpReceiver.PacketNTP(pkt.Timestamp)
}

func (p *Parma) RtpTimeEncode(t time.Time) (uint32, bool) {
	sessionTime, ok := p.SessionTime()
	if !ok {
		return 0, false
	}

	if p.rtpTimeEncoder == nil {
		return 0, false
	}

	return p.rtpTimeEncoder.Encode(t.Sub(sessionTime)), true
}

func (p *Parma) SetSessionTimeFunc(fn SessionTimeFunc) *Parma {
	p._SessionTime = fn
	return p
}

func (p *Parma) OnReadRTP(t Track, packet *rtp.Packet) {
	if p.rtcpReceiver != nil {
		_ = p.rtcpReceiver.ProcessPacket(packet, time.Now(), true)
	}

	if p._OnReadRTP != nil {
		p._OnReadRTP(t, packet)
	}
}
func (p *Parma) SetOnReadRTPFunc(fn OnReadRTPFunc) *Parma {
	p._OnReadRTP = fn
	return p
}

func (p *Parma) OnReadRTCP(t Track, packets []rtcp.Packet) {
	if p.rtcpReceiver != nil {
		for _, packet := range packets {
			switch pkt := packet.(type) {
			case *rtcp.SenderReport:
				p.rtcpReceiver.ProcessSenderReport(pkt, time.Now())
			}
		}
	}

	if p._OnReadRTCP != nil {
		p._OnReadRTCP(t, packets)
	}
}
func (p *Parma) SetOnReadRTCPFunc(fn OnReadRTCPFunc) *Parma {
	p._OnReadRTCP = fn
	return p
}

func (p *Parma) OnFrame(t Track, frame *Frame) {
	if p._OnFrame != nil {
		p._OnFrame(t, frame)
	}
}
func (p *Parma) SetOnFrameFunc(fn OnFrameFunc) *Parma {
	p._OnFrame = fn
	return p
}

func (p *Parma) OnLossRTP(t Track, sequenceNumber uint16, diff int) {
	if p._OnLossRTP != nil {
		p._OnLossRTP(t, sequenceNumber, diff)
	} else {
		log.Printf("[loss rtp] id: %v, ssrc: %v, codec: %v, seq: %v, loss: %v",
			t.ID(), t.SSRC(), t.CtlParma().FmtContext.GetRtpMap().GetCodec(),
			sequenceNumber, diff,
		)
	}
}
func (p *Parma) SetOnLossRTPFunc(fn OnLossRTPFunc) *Parma {
	p._OnLossRTP = fn
	return p
}

func (p *Parma) Detector(t Track, pkt *rtp.Packet) ([]*rtp.Packet, int) {

	packets, n := p.rtpReorder.Process(pkt)
	for _, packet := range packets {
		if diff := p.rtpLossDetector.Process(packet); diff > 0 {
			// 有丢包
			p.OnLossRTP(t, packet.SequenceNumber, diff)
		}
	}

	return packets, n
}
